import json
import os
from typing import List, Dict, Any
from datetime import datetime

class ChatHistoryManager:
    """聊天历史管理器，负责处理对话历史的保存和加载"""
    
    def __init__(self, history_file: str = "chat_history.json"):
        self.history_file = history_file
        self.current_session = []
        self.load_history()
    
    def load_history(self) -> None:
        """加载聊天历史"""
        if os.path.exists(self.history_file):
            try:
                with open(self.history_file, 'r', encoding='utf-8') as f:
                    history = json.load(f)
                    if isinstance(history, list):
                        # 确保每个消息都有正确的 role 和 content
                        valid_messages = []
                        for msg in history:
                            if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                                valid_messages.append(msg)
                        self.current_session = valid_messages[-20:] if valid_messages else []
            except (json.JSONDecodeError, FileNotFoundError):
                self.current_session = []
    
    def save_history(self) -> None:
        """保存聊天历史"""
        try:
            # 如果文件存在，先读取现有历史
            existing_history = []
            if os.path.exists(self.history_file):
                with open(self.history_file, 'r', encoding='utf-8') as f:
                    try:
                        existing_history = json.load(f)
                    except json.JSONDecodeError:
                        pass
            
            # 合并现有历史和当前会话
            if not isinstance(existing_history, list):
                existing_history = []
            
            # 保存当前会话
            existing_history.extend(self.current_session)
            if len(existing_history) > 100:
                existing_history = existing_history[-100:]
            
            with open(self.history_file, 'w', encoding='utf-8') as f:
                json.dump(existing_history, f, indent=4, ensure_ascii=False)
        except Exception as e:
            print(f"保存聊天历史时出错: {str(e)}")
    
    def add_message(self, role: str, content: str) -> None:
        """添加新消息到当前会话"""
        # 确保 role 是有效的
        valid_roles = {"system", "user", "assistant"}
        if role not in valid_roles:
            if role == "AI":
                role = "assistant"
            elif role == "用户":
                role = "user"
            else:
                role = "system"
        
        # 添加消息
        self.current_session.append({"role": role, "content": content})
        if len(self.current_session) > 20:  # 保持当前会话不超过20条消息
            self.current_session = self.current_session[-20:]
        
        # 保存历史记录
        self.save_history()
    
    def clear_current_session(self) -> None:
        """清空当前会话"""
        self.current_session = []
    
    def get_context(self, max_messages: int = 10) -> List[Dict[str, str]]:
        """获取上下文消息"""
        return self.current_session[-max_messages:] if self.current_session else []
    
    def get_all_history(self) -> List[Dict[str, Any]]:
        """获取所有历史记录"""
        all_history = []
        if os.path.exists(self.history_file):
            try:
                with open(self.history_file, 'r', encoding='utf-8') as f:
                    history = json.load(f)
                    if isinstance(history, list):
                        # 按会话分组
                        sessions = []
                        current_session = []
                        session_id = 1
                        
                        for msg in history:
                            if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                                # 添加时间戳（如果没有的话）
                                if 'timestamp' not in msg:
                                    msg['timestamp'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                                current_session.append(msg)
                                
                                # 如果是系统消息或者会话太长，开始新会话
                                if msg['role'] == 'system' and len(current_session) > 1:
                                    if len(current_session) > 1:
                                        sessions.append({
                                            'session_id': session_id,
                                            'messages': current_session[:-1]
                                        })
                                        session_id += 1
                                    current_session = [msg]
                        
                        # 添加最后一个会话
                        if current_session:
                            sessions.append({
                                'session_id': session_id,
                                'messages': current_session
                            })
                        
                        return sessions
            except (json.JSONDecodeError, FileNotFoundError):
                pass
        return []
    
    def load_from_data(self, data: List[Dict[str, Any]]) -> None:
        """从数据加载历史记录"""
        try:
            # 将会话数据转换为消息列表
            messages = []
            for session in data:
                if isinstance(session, dict) and 'messages' in session:
                    messages.extend(session['messages'])
            
            # 保存到文件
            with open(self.history_file, 'w', encoding='utf-8') as f:
                json.dump(messages, f, indent=4, ensure_ascii=False)
            
            # 重新加载历史
            self.load_history()
        except Exception as e:
            print(f"加载历史数据时出错: {str(e)}")
    
    def clear_all_history(self) -> None:
        """清空所有历史记录"""
        try:
            if os.path.exists(self.history_file):
                os.remove(self.history_file)
            self.current_session = []
        except Exception as e:
            print(f"清空历史记录时出错: {str(e)}")